import random
from typing import Tuple, List, Dict, Optional
from chess.board import Board
from chess.pieces import Piece, Pawn, Knight, Bishop, Rook, Queen, King

class ChessAI:
    def __init__(self, color: str = 'black', difficulty: str = 'normal'):
        """
        初始化AI对手
        :param color: AI所执棋子颜色
        :param difficulty: 难度级别 ('easy', 'normal', 'hard')
        """
        self.color = color  # AI所执棋子颜色
        self.difficulty = difficulty  # 难度级别
        # 棋子价值（用于评估局面）
        self.piece_values = {
            Pawn: 1,    # 兵
            Knight: 3,  # 马
            Bishop: 3,  # 象
            Rook: 5,    # 车
            Queen: 9,   # 后
            King: 100   # 王
        }
        # 位置评估表（简化版）
        self.position_values = {
            # 兵的位置评估（鼓励兵向前推进）
            Pawn: [
                [0, 0, 0, 0, 0, 0, 0, 0],
                [5, 5, 5, 5, 5, 5, 5, 5],
                [1, 1, 2, 3, 3, 2, 1, 1],
                [0.5, 0.5, 1, 2.5, 2.5, 1, 0.5, 0.5],
                [0, 0, 0, 2, 2, 0, 0, 0],
                [0.5, -0.5, -1, 0, 0, -1, -0.5, 0.5],
                [0.5, 1, 1, -2, -2, 1, 1, 0.5],
                [0, 0, 0, 0, 0, 0, 0, 0]
            ],
            # 马的位置评估（鼓励马在中心区域）
            Knight: [
                [-5, -4, -3, -3, -3, -3, -4, -5],
                [-4, -2, 0, 0, 0, 0, -2, -4],
                [-3, 0, 1, 1.5, 1.5, 1, 0, -3],
                [-3, 0.5, 1.5, 2, 2, 1.5, 0.5, -3],
                [-3, 0, 1.5, 2, 2, 1.5, 0, -3],
                [-3, 0.5, 1, 1.5, 1.5, 1, 0.5, -3],
                [-4, -2, 0, 0.5, 0.5, 0, -2, -4],
                [-5, -4, -3, -3, -3, -3, -4, -5]
            ],
            # 象的位置评估
            Bishop: [
                [-2, -1, -1, -1, -1, -1, -1, -2],
                [-1, 0, 0, 0, 0, 0, 0, -1],
                [-1, 0, 0.5, 1, 1, 0.5, 0, -1],
                [-1, 0.5, 0.5, 1, 1, 0.5, 0.5, -1],
                [-1, 0, 1, 1, 1, 1, 0, -1],
                [-1, 1, 1, 1, 1, 1, 1, -1],
                [-1, 0.5, 0, 0, 0, 0, 0.5, -1],
                [-2, -1, -1, -1, -1, -1, -1, -2]
            ],
            # 车的位置评估
            Rook: [
                [0, 0, 0, 0, 0, 0, 0, 0],
                [0.5, 1, 1, 1, 1, 1, 1, 0.5],
                [-0.5, 0, 0, 0, 0, 0, 0, -0.5],
                [-0.5, 0, 0, 0, 0, 0, 0, -0.5],
                [-0.5, 0, 0, 0, 0, 0, 0, -0.5],
                [-0.5, 0, 0, 0, 0, 0, 0, -0.5],
                [-0.5, 0, 0, 0, 0, 0, 0, -0.5],
                [0, 0, 0, 0.5, 0.5, 0, 0, 0]
            ],
            # 皇后的位置评估
            Queen: [
                [-2, -1, -1, -0.5, -0.5, -1, -1, -2],
                [-1, 0, 0, 0, 0, 0, 0, -1],
                [-1, 0, 0.5, 0.5, 0.5, 0.5, 0, -1],
                [-0.5, 0, 0.5, 0.5, 0.5, 0.5, 0, -0.5],
                [0, 0, 0.5, 0.5, 0.5, 0.5, 0, -0.5],
                [-1, 0.5, 0.5, 0.5, 0.5, 0.5, 0, -1],
                [-1, 0, 0.5, 0, 0, 0, 0, -1],
                [-2, -1, -1, -0.5, -0.5, -1, -1, -2]
            ],
            # 国王的位置评估（开局和中局，鼓励国王安全）
            King: [
                [-3, -4, -4, -5, -5, -4, -4, -3],
                [-3, -4, -4, -5, -5, -4, -4, -3],
                [-3, -4, -4, -5, -5, -4, -4, -3],
                [-3, -4, -4, -5, -5, -4, -4, -3],
                [-2, -3, -3, -4, -4, -3, -3, -2],
                [-1, -2, -2, -2, -2, -2, -2, -1],
                [2, 2, 0, 0, 0, 0, 2, 2],
                [2, 3, 1, 0, 0, 1, 3, 2]
            ]
        }
    
    def choose_move(self, game) -> Tuple[Tuple[int, int], Tuple[int, int]]:
        """
        选择一个移动
        :param game: 游戏对象
        :return: 移动的起点和终点坐标
        """
        if self.difficulty == 'easy':
            return self._choose_random_move(game)
        elif self.difficulty == 'hard':
            return self._choose_best_move(game, depth=3)
        else:  # normal
            return self._choose_best_move(game, depth=2)
    
    def _choose_random_move(self, game) -> Tuple[Tuple[int, int], Tuple[int, int]]:
        """
        随机选择一个合法移动（简单难度）
        :param game: 游戏对象
        :return: 移动的起点和终点坐标
        """
        valid_moves = self._get_all_valid_moves(game)
        if not valid_moves:
            return None
        
        return random.choice(valid_moves)
    
    def _choose_best_move(self, game, depth: int = 2) -> Tuple[Tuple[int, int], Tuple[int, int]]:
        """
        选择最佳移动（使用极小化极大算法）
        :param game: 游戏对象
        :param depth: 搜索深度
        :return: 最佳移动的起点和终点坐标
        """
        valid_moves = self._get_all_valid_moves(game)
        if not valid_moves:
            return None
        
        best_score = float('-inf')
        best_move = None
        
        for from_pos, to_pos in valid_moves:
            # 模拟移动
            piece = game.board.get_piece_at(from_pos)
            captured_piece = game.board.get_piece_at(to_pos)
            
            # 临时移动
            original_position = piece.position
            original_has_moved = piece.has_moved
            
            game.board.remove_piece_at(from_pos)
            game.board.remove_piece_at(to_pos)  # 移除可能存在的被吃子
            piece.position = to_pos
            piece.has_moved = True
            game.board.place_piece(piece)
            
            # 评估局面
            score = -self._minimax(game, depth - 1, float('-inf'), float('inf'), False)
            
            # 恢复棋盘状态
            game.board.remove_piece_at(to_pos)
            piece.position = original_position
            piece.has_moved = original_has_moved
            game.board.place_piece(piece)
            
            if captured_piece:
                game.board.place_piece(captured_piece)
            
            if score > best_score:
                best_score = score
                best_move = (from_pos, to_pos)
        
        return best_move
    
    def _minimax(self, game, depth: int, alpha: float, beta: float, is_maximizing: bool) -> float:
        """
        极小化极大算法（带Alpha-Beta剪枝）
        :param game: 游戏对象
        :param depth: 当前深度
        :param alpha: Alpha值
        :param beta: Beta值
        :param is_maximizing: 是否为极大化层
        :return: 局面评分
        """
        if depth == 0:
            return self._evaluate_board(game.board)
        
        # 获取有效移动
        current_color = self.color if is_maximizing else ('white' if self.color == 'black' else 'black')
        valid_moves = []
        
        for piece in game.board.get_all_pieces(current_color):
            if game.select_piece(piece.position):
                for move in game.valid_moves:
                    valid_moves.append((piece.position, move))
                
                # 重置选择状态
                game.selected_piece = None
                game.valid_moves = []
        
        if not valid_moves:
            # 如果没有合法移动，检查是否被将军（将死或逼和）
            is_check = game.board.is_king_in_check(current_color)
            if is_check:
                return float('-inf') if is_maximizing else float('inf')  # 被将死
            return 0  # 和棋
        
        if is_maximizing:
            max_eval = float('-inf')
            for from_pos, to_pos in valid_moves:
                # 模拟移动
                piece = game.board.get_piece_at(from_pos)
                captured_piece = game.board.get_piece_at(to_pos)
                
                original_position = piece.position
                original_has_moved = piece.has_moved
                
                game.board.remove_piece_at(from_pos)
                game.board.remove_piece_at(to_pos)
                piece.position = to_pos
                piece.has_moved = True
                game.board.place_piece(piece)
                
                # 递归评估
                eval = self._minimax(game, depth - 1, alpha, beta, False)
                max_eval = max(max_eval, eval)
                alpha = max(alpha, eval)
                
                # 恢复棋盘状态
                game.board.remove_piece_at(to_pos)
                piece.position = original_position
                piece.has_moved = original_has_moved
                game.board.place_piece(piece)
                
                if captured_piece:
                    game.board.place_piece(captured_piece)
                
                if beta <= alpha:
                    break  # Alpha剪枝
            
            return max_eval
        else:
            min_eval = float('inf')
            for from_pos, to_pos in valid_moves:
                # 模拟移动
                piece = game.board.get_piece_at(from_pos)
                captured_piece = game.board.get_piece_at(to_pos)
                
                original_position = piece.position
                original_has_moved = piece.has_moved
                
                game.board.remove_piece_at(from_pos)
                game.board.remove_piece_at(to_pos)
                piece.position = to_pos
                piece.has_moved = True
                game.board.place_piece(piece)
                
                # 递归评估
                eval = self._minimax(game, depth - 1, alpha, beta, True)
                min_eval = min(min_eval, eval)
                beta = min(beta, eval)
                
                # 恢复棋盘状态
                game.board.remove_piece_at(to_pos)
                piece.position = original_position
                piece.has_moved = original_has_moved
                game.board.place_piece(piece)
                
                if captured_piece:
                    game.board.place_piece(captured_piece)
                
                if beta <= alpha:
                    break  # Beta剪枝
            
            return min_eval
    
    def _get_all_valid_moves(self, game) -> List[Tuple[Tuple[int, int], Tuple[int, int]]]:
        """
        获取所有有效移动
        :param game: 游戏对象
        :return: 有效移动列表，每个元素为(起点坐标, 终点坐标)
        """
        valid_moves = []
        
        for piece in game.board.get_all_pieces(self.color):
            if game.select_piece(piece.position):
                for move in game.valid_moves:
                    valid_moves.append((piece.position, move))
                
                # 重置选择状态
                game.selected_piece = None
                game.valid_moves = []
        
        return valid_moves
    
    def _evaluate_board(self, board: Board) -> float:
        """
        评估当前局面
        :param board: 棋盘对象
        :return: 局面评分（正值有利于AI，负值不利于AI）
        """
        score = 0
        
        # 计算双方棋子的总价值
        for piece in board.get_all_pieces():
            piece_value = self.piece_values.get(type(piece), 0)
            position_value = self._get_position_value(piece)
            
            value = piece_value + position_value
            
            if piece.color == self.color:
                score += value
            else:
                score -= value
        
        # 检查将军状态
        opponent_color = 'white' if self.color == 'black' else 'black'
        if board.is_king_in_check(opponent_color):
            score += 0.5  # 将对方将军加分
        if board.is_king_in_check(self.color):
            score -= 0.5  # 被对方将军减分
        
        return score
    
    def _get_position_value(self, piece: Piece) -> float:
        """
        获取棋子在特定位置的额外价值
        :param piece: 棋子对象
        :return: 位置价值分数
        """
        piece_type = type(piece)
        if piece_type not in self.position_values:
            return 0
        
        row, col = piece.position
        
        # 对于黑方棋子，需要翻转棋盘
        if piece.color == 'black':
            row = 7 - row
        
        return self.position_values[piece_type][row][col]
    
    def set_difficulty(self, difficulty: str):
        """
        设置AI难度
        :param difficulty: 难度级别 ('easy', 'normal', 'hard')
        """
        if difficulty in ['easy', 'normal', 'hard']:
            self.difficulty = difficulty 